# pytombo - Python library for processing plain text + encrypted notes.
# Copyright (C) 2007  Chris Clark

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""
chi I/O module

Reads Writes encrypted Tombo *.chi

NOTE: if you want to read notes in Tombo that where generated with
this module, ensure to send in strings to write_encrypted_file()
with Windows style newlines (i.e. '\x0D\x0A')

Tombo is available from http://tombo.sourceforge.jp/En/
"""

"""
TODO list

*   Support file-like objects interface, not just "here is the string"/"return the string"
*   clean up asserts - hey can be "compiled" out in non-debug mode
*   Clean up exception (classes, handle repr) and raising (match pep-8)
*   Run through PyChecker, PyLint, etc. checkers
*   Try and use a consistent coding style:
    http://www.python.org/dev/peps/pep-0008/
    http://www.python.org/dev/peps/pep-0257/
    This is slightly tricky as pep 8 encourages "lowercase, with words separated by underscores" over mixedCase
    May need to rename module too as underscores are frowned upon, e.g. chi_io to chiio which doesn't read well.
*   See http://wiki.python.org/moin/PythonSpeed
*   See http://wiki.python.org/moin/PythonSpeed/PerformanceTips -- contain speed and profiling information/tips
*   UnitTest to generate some data and then crypt/decrpt it. Ideally make use of TomboCrypt for confirmed sanity checks.
*   Profile/time so that if any performance tweaks are made they can be checked! Test with Psyco http://psyco.sourceforge.net/. More struct usage, etc.
    *   most time spent in __round_func() - probably not much can be done
        there, however init of Blowfish is nearly 50% of a small decryption.
        If the same password is used for reading many files this is a large
        overhead that could be reduced to 1 call!
*   Some list, ord(), chr(), to string operations could be simplified to perform more operations whilst as lists before conversion back to strings
    *   use array module (array.array) instead of lists for performance
    *   Remove string operations, try and use cStringIO library instead to save on garbage collection and creating new items
"""

import sys
import struct
import array
try:
    import hashlib
    md5checksum=hashlib.md5
except ImportError:
    # pre 2.6/2.5
    import md5
    md5checksum=md5.new
import string
import random
try:
    #raise ImportError
    import cStringIO as StringIO
except ImportError:
    import StringIO


"""
Import pure python blowfish implementation
this is from http://cheeseshop.python.org/pypi/pypwsafe/0.0.2
this differs from the original Michael Gilfix <mgilfix@eecs.tufts.edu>
version in that it has been up-dated to deal with Long Integers so you do
not get future warnings with Python 2.3 (and based on my experience with
freddb cdkeys) and errors/data-corruption in Python 2.4).

See:
    http://jason.diamond.name/weblog/2005/04/07/cracking-my-password-safe
    http://jason.diamond.name/weblog/2005/10/04/pypwsafe-release-1
    http://jason.diamond.name/weblog/2005/10/05/pypwsafe-0-0-2-with-setup-dot-py
"""
try:
   # raise ImportError
    # Try fast blowfish first
    # http://www.dlitz.net/software/pycrypto/
    from Crypto.Cipher import Blowfish
    TheBlowfishCons = Blowfish.new
    TheBlowfishClass = type(TheBlowfishCons('x'))
    print 'using PyCrypto'
except:
    import blowfish
    TheBlowfishClass = blowfish.Blowfish
    TheBlowfishCons = blowfish.Blowfish


def dump_bytes(s):
    """
    Poor mans hexdump (to stdio)
    """
    mycount = 0
    for c in s:
        print '%02X' % ord(c),
        mycount  = mycount + 1
        if mycount == 16 :
            mycount = 0
            print
    print


def swap_bytes(s):
    """
    Simple swap byte routine - Not actually used, here just in case
    """
    a = array.array('L', s)
    a.byteswap()
    return a.tostring()

class ChiIO(Exception):
    '''Base chi I/O exception'''

class BadPassword(ChiIO):
    '''Bad password exception'''

class UnsupportedFile(ChiIO):
    '''File not encrypted/not supported exception'''

def gen_random_string(length_of_str):
    """generate a srting containing random characters of length length_of_str
    """
    source_set = string.ascii_letters+string.digits+string.punctuation
    result=[]
    for x in range(length_of_str):
        result.append(random.choice(source_set))
    return ''.join(result)

def CHI_cipher(password):
    if isinstance(password, TheBlowfishClass):
        cipher = password
    else:
        # Generate md5 sum of password, this is what is used as the encrypt key
        m = md5checksum()
        m.update(password)
        md5key = m.digest()
        cipher = TheBlowfishCons(md5key)
    return cipher

def read_encrypted_file(fileinfo, password):
    """Reads a *.chi file encrypted by Tombo. Returns (8 bit) string containing plaintext.
    Raises exceptions on failure.
    
    fileinfo is either a filename (string) or a file-like object that reads binary bytes that be can read (caller is responsible for closing)
    password is a (byte) string, i.e. not Unicode type
    """
    if password is None:
        raise BadPassword('None passed in for password for file %r' % fileinfo)
    
    if isinstance(fileinfo, basestring):
        enc_filename=fileinfo
        in_file = open(enc_filename, 'rb')
    else:
        #assume it is a file-like object
        in_file = fileinfo
        ## TODO look up filename from object, file-likes usually have an attribute
        #enc_filename = in_file.name
        enc_filename = None
    
    # called version in CryptManager
    header = in_file.read(4)
    if header != 'BF01':
        raise UnsupportedFile, 'not a Tombo *.chi file'
    
    # read in 4 bytes and convert into an integer value
    ## NOTE may need to worry about byte swap on big-endian hardware
    # TODO do we need array lookup if we do not intend to byte swap???
    tmpbuf = in_file.read(4)
    #dump_bytes(tmpbuf)
    xx = array.array('L', tmpbuf)
    (enc_len,) = struct.unpack('L', xx)

    enc_data = in_file.read()
    encbuf_len = len(enc_data)
    #print 'read in %d bytes, of that only %d byte(s) are real data' % (encbuf_len , enc_len)

    # Assume it is a plain text string (i.e. a byte string, not Unicode type)
    cipher = CHI_cipher(password)
    
    mycounter = encbuf_len
    decrypted_data=[]
    second_pass = list("BLOWFISH")
    while mycounter >= 8:
        data = enc_data[:8]
        chipher=data
        enc_data = enc_data[8:]
        data = cipher.decrypt(data)
        ## based on debug code (and tombo specific additions to blowfish.c) in Tombo
        ## tombo is using the base blowfish alogrithm AND then applies more bit fiddling....
        ## performs bitwise exclusive or on decrypted text from blowfish and "BLOWFISH" (note this static gets modified....)
        for x in range(8):
            decrypted_data.append( ord(data[x]) ^ ord(second_pass[x]) )
            second_pass[x] = chipher[x]
        mycounter = mycounter - 8
    ## there should be no bytes left after this.
    ## Ignore any remaining bytes (less than 8)?
    if mycounter > 0:
        # This should not happen if it did this may be a corrupted file
        # at present not handled, pending on bugs found here
        raise RuntimeError, 'ExtraBytesFound during decryption'
    decrypted_data = map(chr, decrypted_data)
    decrypted_data = ''.join(decrypted_data)
    """
    At this point decrypted_data contains:
        8 bytes of (unknown) random data
        16 bytes of an md5sum of the unencrypted data
        enc_len * bytes of unencrypted data
    
    But at this point we do not know if the password that was used was correct,
    the unencrypted data could be garbage!
    """
    
    # extract real text from supuriuos crap (24 bytes on from started of data)
    # loose spurious end use real data length we read earlier?
    unencrypted_str = decrypted_data[24:]
    unencrypted_str = unencrypted_str[:enc_len]
    
    m = md5checksum()
    m.update(unencrypted_str)
    decriptsum = m.digest()
    chi_md5sum = decrypted_data[8:]
    chi_md5sum = chi_md5sum[:16]
    
    if chi_md5sum == decriptsum:
        # passwords match, so data is valid
        return unencrypted_str
    else:
        #password did not match, data is bogus
        # raise exception WITH information such as filename, do not dump out password as that could be a security hole
        raise BadPassword, enc_filename

def write_encrypted_file(fileinfo, password, plaintext):
    """Writes an encrypted *.chi file that could be read by Tombo. Parameter plaintext should be 8 bit string.
    Raises exceptions on failure (so caller is responsible for cleaning up incomplete out files).
    NOTE: if notes created with this routine are to be read in Tombo
    ensure to send in plaintext strings with Windows style newlines;
    i.e. '\x0D\x0A'. See dumb_unix2dos().
    
    fileinfo is either a filename (string) or a file-like object that writes binary bytes that be can written to (caller is responsible for closing)
    password is a (byte) string, i.e. not Unicode type
    
    """
    assert isinstance(plaintext, str), 'Only support 8 bit plaintext. Encode first, see help(codecs).'
    
    plain_text = plaintext
    if isinstance(fileinfo, basestring):
        enc_filename=fileinfo
    else:
        #assume it is a file-like object
        enc_filename = None
    
    plain_text_len = len(plain_text)
    #print 'read in %d bytes' % plain_text_len
    
    m = md5checksum()
    m.update(plain_text)
    plain_text_md5sum = m.digest()
    
    #print "plain_text_md5sum"
    #dump_bytes(plain_text_md5sum)
    
    # Generate md5 sum of password, this is what is used as the encrypt key
    cipher = CHI_cipher(password)
    
    ## create new buffer that will be encrypted, contains:
    ##  8 bytes of random (if this is NOT random, then a plaintext+password will always create the SAME encrypted text)
    ##  16 bytes of md5 of plaintext
    ##  plain_text_len bytes of plaintext
    #str_to_encrypt = '12345678' + plain_text_md5sum + plain_text
    #enc_data = '12345678' + plain_text_md5sum + plain_text
    enc_data = gen_random_string(8) + plain_text_md5sum + plain_text
    
    mycounter = len(enc_data)
    encrypted_data=''
    second_pass = "BLOWFISH"
    while mycounter >= 8:
        data = enc_data[:8]
        
        ## based on debug code (and tombo specific additions to blowfish.c) in Tombo
        ## tombo is using the base blowfish alogrithm AND then applies more bit fiddling....
        ## performs bitwise exclusive or on decrypted text from blowfish and "BLOWFISH" (note this static gets modified....)
        plain=[]
        for x in range(8):
            plain.append( ord(data[x]) ^ ord(second_pass[x]) )
        #print 'dec pass2 ', plain
        data = map(chr, plain)
        data = ''.join(data )
        #print data
        #dump_bytes(data)
        
        
        enc_data = enc_data[8:]
        ##debug
        #first_byte = ord(data[0])
        #print 'raw : %x, %d' %(first_byte , first_byte )
        #dump_bytes(data)
        data = cipher.encrypt(data)
        
        # take encrypted block and shove into second pass but fiddler (really first pass when encrypting
        second_pass=[]
        for x in range(8):
            second_pass.append(data[x])
        second_pass = ''.join(second_pass)
    
        ##debug
        #first_byte = ord(data[0])
        #print 'encrypted : %x, %d' %(first_byte , first_byte )
        #print ""
        encrypted_data = encrypted_data + data
        mycounter = mycounter - 8
    ## there maybe a few odd bytes left after this (less than 8). Take then and put with random bytes into an 8 byte block and encrypt that
    if mycounter > 0:
        # take  last few bytes "data"
        tmp_buf=[]
        for x in range(mycounter):
            tmp_buf.append(enc_data[x])
        data = ''.join(tmp_buf)
        
        # Tombo bit fiddling 
        # NOTE in Tobo the "fake" padding bytes at the end are not bit fiddled with 
        plain=[]
        for x in range(mycounter):
            plain.append( ord(data[x]) ^ ord(second_pass[x]) )

        data = map(chr, plain)
        
        # pad the end few bytes so that blowfish can be applied
        # just take "garbage" from begnning of (last/previously) encrypted block
        # NOTE this differs from Tombo which takes the garbage from the end of the previously encrypted block
        for x in range(8 - mycounter):
            data.append(data[x])

        data = ''.join(data )

        # Now encrypt
        data = cipher.encrypt(data)
        encrypted_data = encrypted_data + data
    
    ## file IO section
    if enc_filename is not None:
        out_file = open(enc_filename, 'wb')
    else:
        #assume it is a file-like object
        out_file = fileinfo
    
    # called version in tombo's CryptManager
    header = 'BF01'
    
    out_file.write(header)
    
    # write out 4 bytes - integer value containing unencrypted length of data
    # write out plaintext length
    out_file.write(struct.pack('L', plain_text_len))
    
    # write out encrypted data
    out_file.write(encrypted_data)
    
    if enc_filename is not None:
        # i.e. we opened the file so we need to close it
        out_file.close()


def dumb_unix2dos(in_str):
    """In-efficient but simple unix2dos string conversion
    convert '\x0A' --> '\x0D\x0A'
    """
    return in_str.replace('\x0A', '\x0D\x0A')


## Consider using filelike from http://cheeseshop.python.org/pypi/filelike/
class ChiAsFile(object):
    """does not really honor seek(), etc. only read/write"""
    def __init__(self, fileptr , password):
        self._fileptr = fileptr
        self._password = password
        self._bufferedfileptr = StringIO.StringIO()
        self._fileread = None
    
    def __getattr__(self, attr):
        if self.__dict__.has_key(attr):
            return self.__dict__[attr]
        else:
            return getattr(self._bufferedfileptr, attr)
    
    def read(self, size=None):
        if self._fileread is None:
            self._fileread = True
            plain_text = read_encrypted_file(self._fileptr, self._password)
            self._bufferedfileptr = StringIO.StringIO(plain_text)
        if not self._fileread:
            raise IOError('file was write, and then read issued. read and write are mutually exclusive operations')
        # do we need this if?
        if size is None:
            return self._bufferedfileptr.read()
        else:
            return self._bufferedfileptr.read(size)
            
    def write(self, str_of_bytes):
        if self._fileread is None:
            self._fileread = False
        if self._fileread:
            raise IOError('file was read, and then write issued. read and write are mutually exclusive operations')
        return self._bufferedfileptr.write(str_of_bytes)
        
    def close(self, *args, **kwargs):
        ## do we need to call this in __del__?
        if not self._fileread:
            # i.e writable file
            plain_text = self._bufferedfileptr.getvalue()
            write_encrypted_file(self._fileptr, self._password, plain_text)
        #self._fileptr.close() # is this right?
        self._bufferedfileptr.close()
        ## TODO disallow more writes/closes....


def demo_test():
    """Examples of use with some error handling
    """
    print "tombo test example"
    
    plain_text = """12345678"""
    enc_fname = 'chi_io_test1.chi'
    mypassword = 'testing'
    
    print 'Write encrypted file to', enc_fname
    write_encrypted_file(enc_fname, mypassword, plain_text)
    try:
        print 'Reading encrypted file', enc_fname
        plain_str = read_encrypted_file(enc_fname, mypassword)
        print '--------- Decrypted data ---------'
        print plain_str 
        print '-------------- End ---------------'
    except BadPassword, info:
        print "bad password used"
        print info
    except UnsupportedFile, info:
        print "file was not encrypted or is not supported"
        print info
    print '\n***************************\n'
        
    
    plain_text = """Line 1
of a simple text file.
"""
    enc_fname = 'chi_io_test.chi'
    mypassword = 'testing'
    
    print 'Write multiline encrypted file to', enc_fname
    write_encrypted_file(enc_fname, mypassword, plain_text)
    try:
        print 'Reading encrypted file', enc_fname
        plain_str = read_encrypted_file(enc_fname, mypassword)
        print '--------- Decrypted data ---------'
        print plain_str 
        print '-------------- End ---------------'
    except BadPassword, info:
        print "bad password used"
        print info
    except UnsupportedFile, info:
        print "file was not encrypted or is not supported"
        print info
    print '\n***************************\n'
    mypassword = 'bad password'
    try:
        print 'Reading encrypted file with bad password', enc_fname
        plain_str = read_encrypted_file(enc_fname, mypassword)
        print '--------- Decrypted data ---------'
        print plain_str 
        print '-------------- End ---------------'
    except BadPassword, info:
        print "bad password used"
        print info
    except UnsupportedFile, info:
        print "file was not encrypted or is not supported"
        print info
    print '\n***************************\n'
    
    
    enc_fname = 'delete_me.txt'
    mypassword='testing'
    myfile = open(enc_fname, 'w')
    myfile.write('hello world')
    myfile.close()
    try:
        print 'Reading non-encrypted file', enc_fname
        plain_str = read_encrypted_file(enc_fname, mypassword)
        print '--------- Decrypted data ---------'
        print plain_str 
        print '-------------- End ---------------'
    except BadPassword, info:
        print "bad password used"
        print info
    except UnsupportedFile, info:
        print "file was not encrypted or is not supported"
        print info
    print '\n***************************\n'
    
    enc_fname = 'doesnotexist.chi'
    mypassword='testtest'
    try:
        print 'Reading Non-existent encrypted file', enc_fname
        plain_str = read_encrypted_file(enc_fname, mypassword)
        print '--------- Decrypted data ---------'
        print plain_str 
        print '-------------- End ---------------'
    except BadPassword, info:
        print "bad password used"
        print info
    except UnsupportedFile, info:
        print "file was not encrypted or is not supported"
        print info
    except IOError, info:
        print "Some kind of IO error, probably a missing file"
        print info
    print '\n***************************\n'
    
    plain_text = unicode(u'123\N{POUND SIGN}')
    enc_fname = 'unicode.chi'
    mypassword = 'testing'
    
    try:
        print 'Write encrypted file to', enc_fname
        write_encrypted_file(enc_fname, mypassword, plain_text)
    except AssertionError, info:
        print "assert error, probably tried to feed in Unicode"
        print info
    print '\n***************************\n'
    
    
    plain_text = unicode(u'123\N{POUND SIGN}')
    enc_fname = 'unicode.chi'
    mypassword = 'testing'
    
    try:
        print 'Write encrypted file with unicode to', enc_fname
        
        plain_text = plain_text.encode('utf-8') ## convert to binary string, could be UTF-16, etc
        write_encrypted_file(enc_fname, mypassword, plain_text)
        
        plain_str = read_encrypted_file(enc_fname, mypassword)
        plain_str  = plain_str.decode('utf-8') ## convert to Unicode string
        print '--------- Decrypted data ---------'
        print plain_str 
        print '-------------- End ---------------'
    except AssertionError, info:
        print "assert error, probably tried to feed in Unicode"
        print info
    print '\n***************************\n'
    
    import os
    os.remove('chi_io_test.chi')
    os.remove('chi_io_test1.chi')
    os.remove('unicode.chi')
    
    return 0

def main(argv=None):
    if argv is None:
        argv = sys.argv
    
    print 'argv', argv
    if len(argv) >= 2:
        print 'argv[1]', argv[1]
    
    do_what = raw_input('Encrypt (e) or Decrupt (d), anything else run demo test?: ')
    
    print "do_what", do_what, len(do_what)
    
    if do_what != 'e' and do_what != 'd':
        return demo_test()
        
    enc_fname = raw_input('enter in encrypted filename (e.g. test.chi): ')
    
    import getpass
    mypassword = getpass.getpass("Password:")
    
    if do_what == 'e':
        plain_fname = raw_input('enter in plain text filename (e.g. test.txt): ')
        in_file = open(plain_fname, 'rb')
        plain_text = in_file.read()
        in_file.close()
        print 'Write encrypted file to', enc_fname
        write_encrypted_file(enc_fname, mypassword, plain_text)
    else: ## do_what == 'd':
        try:
            print 'Reading encrypted file', enc_fname
            plain_str = read_encrypted_file(enc_fname, mypassword)
            print '--------- Decrypted data ---------'
            print plain_str 
            print '-------------- End ---------------'
        except BadPassword, info:
            print "bad password used"
            print info
        except UnsupportedFile, info:
            print "file was not encrypted or is not supported"
            print info
    return 0
        

if __name__ == "__main__":
    sys.exit(main())
